
import os
from skimage import io
from skimage.color import rgb2gray
import torch
import matplotlib.pyplot as plt

def show_img(img_t):
    plt.figure()
    for i in range(len(img_t)):
        img_item = img_t[i]
        title = img_item.get('title')
        img = img_item.get('img')

        plt.subplot(len(img_t), 1, i+1)
        plt.suptitle(title)
        plt.imshow(img, cmap="gray")
    plt.show()

def get_img(img):
    read_img = io.imread(img)
    img_t = torch.from_numpy(read_img)
    return img_t

def get_gray_img(img):
    ''' 灰度 '''
    # as_gray=True 代表读出的灰度图的np对象，然后转换成torch的tensor对象
    read_img = io.imread(img)
    img = rgb2gray(read_img)
    img_t = torch.from_numpy(img)
    return img_t

def show_imgtensor(img_t):
    print(img_t)
    print("shape:{}".format(img_t.shape))
    img_t_flat = img_t.view(1, -1)
    print("max:{}".format(img_t_flat.max(dim=1)))
    print("min:{}".format(img_t_flat.min(dim=1)))

def binarization(img_t):
    # 二值化
    z = torch.zeros(30, 120).double()
    o = torch.ones(30, 120).double()
    img_t_b = torch.where(img_t < 0.4, z, o)
    return img_t_b

def spilt(t, label, margin, slice=4):
    ls_image = []
    ls_lab = []
    for i in range(slice):
        s_p = t[0 + margin[0]:30 - margin[1], 30 * i + margin[2]:30 * (i + 1) - margin[3]]
        ls_image.append(s_p)
        ls_lab.append(label[i])
    return ls_image, ls_lab

def show_images(images,labels):
    plt.figure()
    for i in range(len(images)):
        plt.subplot(1, len(images), i+1)
        plt.suptitle(labels)
        plt.imshow(images[i])
    plt.show()

def save_train(base_path):
    # cwd = os.getcwd()
    markimg_path = os.path.join(base_path, 'markimg')
    fs = os.listdir(markimg_path)
    train_path = os.path.join(base_path, 'img_train')

    for f in fs:
        image_path = os.path.join(markimg_path, f)
        img_t = get_gray_img(image_path)
        img_t_b = binarization(img_t)
        f_name,pre = os.path.splitext(f)

        images, labels = spilt(img_t_b, f_name, margin=[2, 2, 2, 1])

        for index in range(len(list(zip(images, labels)))):
            image = images[index]
            label = labels[index]
            # print(label)
            lab_dir = os.path.join(train_path, label)
            if(not os.path.exists(lab_dir)):
                os.mkdir(lab_dir)
            
            image = (image * 255).type(dtype=torch.uint8)
            # 将2a45中的a图片保存成形如2a45_2.png的形式，方便以后追述
            new_name = f"{f_name}_{index}{pre}"
            image_f = os.path.join(lab_dir, new_name)
            print(f'image_f:{image_f}')
            io.imsave(image_f, image)


def main():
    img_paths = ["./markimg/0AVJ.jpeg",
    "./markimg/0CWT.jpeg",
    "./markimg/0D2H.jpeg",
    "./markimg/0EH7.jpeg",
    "./markimg/0FTP.jpeg"]

    img_path = img_paths[3]

    # 原始图片
    src_img = get_img(img_path)

    # 灰度图片
    gray_img = get_gray_img(img_path)
    
    # 二值化
    bin_img = binarization(gray_img)

    images, labels = spilt(bin_img, '0CWT', margin=[2, 2, 2, 1])

    # show_img([{
    #     'title': '原始图片',
    #     'img': src_img
    # },{
    #     'title': '灰度图片',
    #     'img': gray_img
    # },{
    #     'title': '二值化',
    #     'img': bin_img
    # }])

    show_images(images, labels)
    # show_img(img_t)


if __name__ == "__main__":
    save_train('./learn')
    # save_train('./image_test')
